#  -*- coding:utf-8
'''

@File : m3u8DecryptHandler

@Author: syl

@Date : 2020-05-08

@Desc :
1.from Crypto.Cipher import AES
python 在 Windows下使用AES时要安装的是pycryptodome 模块
python 在 Linux下使用AES时要安装的是pycrypto模块
2.已下载ts文件 ，且已有key （文件或bytes）

'''
import json
import os
import shutil
import time
import m3u8
import requests
from glob import iglob
from natsort import natsorted
from Crypto.Cipher import AES
from concurrent.futures import ThreadPoolExecutor


def get_ts_url(m3u8_url):
    m3u8_obj = m3u8.load(m3u8_url)
    base_uri = m3u8_obj.base_uri

    for seg in m3u8_obj.segments:
        yield base_uri + seg.uri


def download_ts(urlinfo):
    url, ts_name = urlinfo
    headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/67.0.3396.99 Safari/537.36', }
    Res = requests.get(url, headers=headers)
    with open(ts_name, 'wb') as fp:
        fp.write(Res.content)


def decrypt_single_ts_keyfile(index, file_name, keyfile, iv=None):
    key = open(keyfile, 'rb').read()
    decrypt_single_ts_key(index, file_name, key, iv)


def decrypt_single_ts_key(index, file_name, key, iv=None):
    data = open(file_name, 'rb').read()
    cipher = AES.new(key, AES.MODE_CBC)
    if iv is not None:
        cipher = AES.new(key, AES.MODE_CBC, iv=iv)
    plain_data = cipher.decrypt(data)

    open(f'temp/{index}.ts', 'wb+').write(plain_data)
    print('file: ' + file_name + '\tsucceed!')


def combin_ts(ts_dir=os.path.dirname(__file__)):
    ts_path = os.path.join(ts_dir, 'temp/*.ts')
    new_path = os.path.join(ts_dir, 'temp/new')  # 防止被处理,写判断多余，直接重命名即可
    with open(new_path, 'wb') as fn:
        for ts in natsorted(iglob(ts_path)):
            with open(ts, 'rb') as ft:
                scline = ft.read()
                fn.write(scline)
    for ts in natsorted(iglob(ts_path)):  # 删除下载的ts文件
        os.remove(ts)
    os.rename(new_path, new_path + '.ts')


if __name__ == '__main__':
    ts_dir = os.path.join(os.path.dirname(__file__), 'children')
    ts_path = os.path.join(ts_dir, '*.ts')
    i = 0
    if not os.path.exists(os.path.join(os.path.dirname(__file__), 'temp')):
        os.mkdir(os.path.join(os.path.dirname(__file__), 'temp'))
    with open('index.key', 'rb') as f:
        key = f.read()
        for ts in natsorted(iglob(ts_path)):
            i += 1
            decrypt_single_ts_key(i, ts, key)
            pass
    start = time.time()
    combin_ts()
    end = time.time()
    print(f"用时:{end - start}")
    fn = 'new.ts'
    output = "td.mp4"
    with open('meta.json.txt', 'rb') as f:
        j = json.loads(f.read())
        output = f"{j.get('title', 'td')}.mp4".replace(" ", "")
    print(output)
    ffmpeg = '"D:\\ffmpeg\\ffmpeg.exe"'
    cmd = ffmpeg + " -i temp\\" + fn + " -c copy " + output
    os.system(cmd)
    shutil.rmtree('temp')
    print('done')
